package cn.crsc.openlookeng.adhoc.jdbc;

import cn.crsc.openlookeng.adhoc.dataframe.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.sql.*;
import java.util.ArrayList;
import java.util.List;

@Slf4j
public class JdbcResultSetDataFrameWrapper implements ResultSetDataFrameWrapper<ResultSet> {
    private Statement st = null;
    private PreparedStatement ps = null;
    private Connection connection = null;
    private QueryObject queryObject = null;

    public JdbcResultSetDataFrameWrapper setResources(QueryObject queryObject, Statement st, PreparedStatement ps, Connection connection) {
        this.st = st;
        this.ps = ps;
        this.connection = connection;
        this.queryObject = queryObject;
        return this;
    }

    @Override
    public PagedDataFrame wrapData(ResultSet result) {
        try {
            ResultSetMetaData metaData = result.getMetaData();

            RowSchemaInfo.RowSchemaInfoBuilder rowSchemaInfoBuilder = RowSchemaInfo.newSchemaBuilder();
            int columnCount = metaData.getColumnCount();
            for (int i = 1; i <= columnCount; i++) {
                String columnName = metaData.getColumnName(i);
                DataType dataType = JdbcTypeMapping.getDataType(JDBCType.valueOf(metaData.getColumnType(i)));
                rowSchemaInfoBuilder.column(columnName, dataType);
            }
            RowSchemaInfo rowSchemaInfo = rowSchemaInfoBuilder.build();
            DataFrameImpl dataFrame = new DataFrameImpl(rowSchemaInfo);
            PagedDataFrame pagedDataFrame = fillDataFrame(result, rowSchemaInfo, dataFrame);
            fillCountNum(st,pagedDataFrame);
            pagedDataFrame.setLoading(false);
            return pagedDataFrame;
        } catch (SQLException e) {
            throw new RuntimeException("sql exception" + e.getMessage());
        }finally {
            closeClosables(result, st, ps);
        }

    }

    private void fillCountNum(Statement ps, PagedDataFrame pagedDataFrame) throws SQLException {
        if(queryObject.getPageSize() > 0 && queryObject.getPageNum() > 0){
            String tableName = StringUtils.substringAfterLast(queryObject.getSql(), " ");
            ResultSet resultSet = st.executeQuery("select count(*) from " + tableName);
            if(resultSet.next()){
                ResultSetColumnReader.BigIntReader bigIntReader = new ResultSetColumnReader.BigIntReader();
                Long count = bigIntReader.readValue(1,resultSet);
                pagedDataFrame.setTotalCount(count);
            }
        }
        return;

    }

    private PagedDataFrame fillDataFrame(ResultSet result, RowSchemaInfo rowSchemaInfo, DataFrameImpl dataFrame) throws SQLException {

        List<ResultSetColumnReader> readers = new ArrayList<>(rowSchemaInfo.getColumns().size());
        for (ColumnInfo columnInfo : rowSchemaInfo.getColumns()) {
            readers.add(JdbcTypeMapping.getResultSetColumnReader(columnInfo.getType()));
        }

        int num = 1;
        PagedDataFrame pagedDataFrame = new PagedDataFrame(rowSchemaInfo, dataFrame);
        pagedDataFrame.setNowPageIndex(queryObject.getPageNum());
        pagedDataFrame.setNowPageSize(queryObject.getPageSize());

        int pageSize = queryObject.getPageSize();
        int pageNum = (queryObject.getPageNum() - 1) * pageSize;
        int untilNum = queryObject.getPageNum() * pageSize;
        if(queryObject.getPageNum() <= 0 || queryObject.getPageSize() <= 0){
            untilNum = Integer.MAX_VALUE;
        }

        while (result.next()) {
            if(num > pageNum && num <= untilNum ){
                List<Object> values = new ArrayList<>(readers.size());
                for (int i = 0; i < readers.size(); i++) {
                    Object value = readers.get(i).readValue(i + 1, result);
                    values.add(value);
                }
                dataFrame.append(values);
                pagedDataFrame.setRowCount(num);
            }
            if (++num > untilNum) {
                break;
            }


        }

        return pagedDataFrame;

    }


    private void closeClosables(AutoCloseable... closeables) {
        for (AutoCloseable closeable : closeables) {
            if (closeable != null) {
                try {
                    log.info("close resource :"+ closeable);
                    closeable.close();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }
    }
}
